library(data.table)
library(ggplot2)
library(tidyverse)
library(gridExtra)
library(grid)
library(readr)
library(stringr)
library(latex2exp)

paper <- TRUE


dir.project <- './'
dir.generated <- './data/generated/matlab/'
dir.fig.paper <-
  paste0(dir.project, 'doc/paper/figures/robustness/')

source('./src/r/figures.R')

## vary EPS

EPS_W_M <- read.csv(paste0(dir.generated, 'robustness/EPS_W_M.csv'), header = FALSE) %>% as.matrix()
EPS_W_C <- read.csv(paste0(dir.generated, 'robustness/EPS_W_C.csv'), header = FALSE) %>% as.matrix()
data_eps <- expand.grid(X=EPS_W_M, Y=EPS_W_C)
TAU_B_EPS <- read.csv(paste0(dir.generated, 'robustness/TAU_B_EPS.csv'), header = FALSE) %>% as.vector()
data_eps$Z <- TAU_B_EPS[,1]


tau_plot_eps <- ggplot(data_eps, aes(X, Y),Z) + 
  geom_raster(aes(fill=Z*100), interpolate = TRUE) +
    scale_fill_gradient2(TeX('t_B'),low='blue', mid='white', high='red', midpoint=0) +
    theme_bw(base_size = 8) +
    xlab(TeX('$\\epsilon_{w_M/w_R,K_B}$')) +
    ylab(TeX('$\\epsilon_{w_C/w_R,K_B}$')) +
    theme.serif


ggsave(
  paste0(dir.fig, "plot_simple_formula_eps.pdf"),
  plot = tau_plot_eps,
  width = 12,
  height = 10,
  units = "cm"
)

## vary PSI

PSI_M <- read.csv(paste0(dir.generated, 'robustness/PSI_M.csv'), header = FALSE) %>% as.matrix()
PSI_C <- read.csv(paste0(dir.generated, 'robustness/PSI_C.csv'), header = FALSE) %>% as.matrix()
data_psi <- expand.grid(X=PSI_M, Y=PSI_C)
TAU_B_PSI <- read.csv(paste0(dir.generated, 'robustness/TAU_B_PSI.csv'), header = FALSE) %>% as.vector()
data_psi$Z <- TAU_B_PSI[,1]


tau_plot_psi <- ggplot(data_psi, aes(X, Y),Z) + 
  geom_raster(aes(fill=Z*100), interpolate = TRUE) +
  scale_fill_gradient2(TeX('t_B'),low='blue', mid='white', high='red', midpoint=0) +
  theme_bw(base_size = 8) +
  xlab(TeX('$\\psi_{M}$')) +
  ylab(TeX('$\\psi_{C}$')) +
  theme.serif


ggsave(
  paste0(dir.fig, "plot_simple_formula_psi.pdf"),
  plot = tau_plot_psi,
  width = 12,
  height = 10,
  units = "cm"
)

plots_combined <-
    grid.arrange(
        tau_plot_eps + ggtitle('A: Effect of wage elasticities'),
        tau_plot_psi + ggtitle('B: Effect of welfare weights '),
        ncol = 2,
        nrow = 1
    )

ggsave(
    paste0(dir.fig, "plots_simple_formula_combined.pdf"),
    plot = plots_combined,
    width = 13,
    height = 6,
    units = "cm"
)
